import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

from fastapi import FastAPI, File, UploadFile, HTTPException, Depends, Form
from fastapi.responses import JSONResponse, HTMLResponse
from contextlib import asynccontextmanager
import cv2
import numpy as np
from insightface.app import FaceAnalysis
from pymilvus import MilvusClient, DataType
import uvicorn
from typing import List, Optional, Dict, Any
import logging
from pydantic_settings import BaseSettings
import json
import asyncio
import time
import datetime
import requests
from fastapi.middleware.cors import CORSMiddleware
import pika
import threading
from concurrent.futures import ThreadPoolExecutor
from queue import Queue
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class Settings(BaseSettings):
    milvus_uri: str = "http://39.104.78.210:19530"
    milvus_token: str = "root:Milvus"
    collection_name: str = "faceD"
    face_det_size: tuple = (640, 640)
    similarity_threshold: float = 0.5
    max_history_size: int = 1000
    rabbitmq_host: str = "localhost"  # 修改为实际服务器IP
    rabbitmq_port: int = 5672  # 使用标准RabbitMQ端口
    rabbitmq_user: str = "root"  # 使用管理员账号
    rabbitmq_password: str = "xeon?YES:N0"  # 使用您提供的密码
    rabbitmq_queue: str = "face_detection_history"
    rabbitmq_virtual_host: str = "/"

    class Config:
        env_file = ".env"


class FaceRecognitionService:
    def __init__(self, settings: Settings):
        self.settings = settings
        self.cli = None
        self.face_app = None
        self.detection_history = []
        self.rabbitmq_connection = None
        self.rabbitmq_channel = None
        self.rabbitmq_connected = False
        self.thread_pool = ThreadPoolExecutor(max_workers=2)
        self.frame_queue = Queue(maxsize=10)

    async def connect_milvus(self):
        try:
            self.cli = MilvusClient(
                uri=self.settings.milvus_uri,
                token=self.settings.milvus_token,
            )
            logger.info(f"Milvus连接成功: {self.settings.milvus_uri}")
            return True
        except Exception as e:
            logger.error(f"Milvus连接失败: {str(e)}")
            raise

    def connect_rabbitmq(self):
        """增强版RabbitMQ连接方法，带重试机制"""
        max_retries = 3
        retry_delay = 2

        for attempt in range(max_retries):
            try:
                credentials = pika.PlainCredentials(
                    self.settings.rabbitmq_user,
                    self.settings.rabbitmq_password
                )
                parameters = pika.ConnectionParameters(
                    host=self.settings.rabbitmq_host,
                    port=self.settings.rabbitmq_port,
                    virtual_host=self.settings.rabbitmq_virtual_host,
                    credentials=credentials,
                    heartbeat=600,
                    blocked_connection_timeout=300,
                    socket_timeout=10
                )

                self.rabbitmq_connection = pika.BlockingConnection(parameters)
                self.rabbitmq_channel = self.rabbitmq_connection.channel()

                # 声明队列（持久化）
                self.rabbitmq_channel.queue_declare(
                    queue=self.settings.rabbitmq_queue,
                    durable=True,
                    arguments={
                        'x-message-ttl': 86400000,  # 消息存活时间24小时
                        'x-max-length': 10000  # 限制队列最大长度
                    }
                )

                self.rabbitmq_connected = True
                logger.info(f"RabbitMQ连接成功: {self.settings.rabbitmq_host}:{self.settings.rabbitmq_port}")
                return True

            except pika.exceptions.AMQPConnectionError as e:
                logger.warning(f"RabbitMQ连接失败(尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt < max_retries - 1:
                    time.sleep(retry_delay)
                    continue
                self.rabbitmq_connected = False
                return False
            except Exception as e:
                logger.error(f"意外的RabbitMQ连接错误: {str(e)}")
                self.rabbitmq_connected = False
                return False

    async def init_face_model(self):
        try:
            self.face_app = FaceAnalysis(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
                                         det_thresh=0.3,
                                         det_box_scale=1.2
                                         )
            self.face_app.prepare(ctx_id=0, det_size=self.settings.face_det_size)
            logger.info("人脸分析模型加载成功")
            return True
        except Exception as e:
            logger.error(f"人脸模型初始化失败: {str(e)}")
            raise

    def read_image(self, file: UploadFile) -> np.ndarray:
        try:
            contents = file.file.read()
            nparr = np.frombuffer(contents, np.uint8)
            img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            return img
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"图片读取失败: {str(e)}")

    async def extract_face_embedding(self, image: UploadFile) -> List[float]:
        img = self.read_image(image)
        faces = self.face_app.get(img)

        if not faces:
            raise HTTPException(status_code=400, detail="未检测到人脸")

        if len(faces) > 1:
            raise HTTPException(status_code=400, detail="检测到多个人脸，请上传单人照片")

        return faces[0].embedding.tolist()

    async def add_face_to_db(self, key: str, embedding: List[float]) -> Dict[str, Any]:
        data = [{'key': key, 'vector': embedding}]
        res = self.cli.insert(
            collection_name=self.settings.collection_name,
            data=data
        )
        return res

    async def search_similar_faces(self, embedding: List[float], limit: int = 1) -> List[Dict]:
        res = self.cli.search(
            collection_name=self.settings.collection_name,
            anns_field="vector",
            data=[embedding],
            limit=limit,
            output_fields=["key"],
            search_params={"metric_type": "COSINE"}
        )

        results = []
        for match in res[0]:
            distance = float(match['distance'])
            key = match['entity']['key'] if distance >= self.settings.similarity_threshold else "Unknown"

            results.append({
                "key": key,
                "distance": distance,
                "confidence": 1 - distance
            })

        return results

    def send_to_rabbitmq(self, message: Dict[str, Any]):
        """增强版消息发送方法"""
        if not self.rabbitmq_connected:
            logger.warning("RabbitMQ未连接，消息无法发送")
            return False

        try:
            # 检查连接是否有效
            if not (self.rabbitmq_connection and self.rabbitmq_connection.is_open):
                logger.warning("RabbitMQ连接已断开，尝试重新连接...")
                if not self.connect_rabbitmq():
                    return False

            message_body = json.dumps(message)
            self.rabbitmq_channel.basic_publish(
                exchange='',
                routing_key=self.settings.rabbitmq_queue,
                body=message_body,
                properties=pika.BasicProperties(
                    delivery_mode=2,  # 持久化消息
                    content_type='application/json',
                    timestamp=int(time.time())
                )
            )
            logger.debug(f"消息已发送到RabbitMQ: {message.get('type')}")
            return True
        except Exception as e:
            logger.error(f"发送消息到RabbitMQ失败: {str(e)}")
            return False

    def close_connections(self):
        """关闭所有连接"""
        if self.rabbitmq_connection and self.rabbitmq_connection.is_open:
            self.rabbitmq_connection.close()
            logger.info("RabbitMQ连接已关闭")

        if self.thread_pool:
            self.thread_pool.shutdown(wait=False)


# Service Instance
face_service = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global face_service
    settings = Settings()
    face_service = FaceRecognitionService(settings)

    try:
        await face_service.connect_milvus()
        await face_service.init_face_model()

        # 同步连接RabbitMQ（带重试）
        for i in range(3):
            if face_service.connect_rabbitmq():
                break
            if i < 2:
                logger.warning(f"RabbitMQ连接失败，10秒后重试... (尝试 {i + 1}/3)")
                await asyncio.sleep(10)

        if not face_service.rabbitmq_connected:
            logger.error("RabbitMQ连接最终失败，服务将以降级模式运行")

        logger.info("人脸识别服务启动完成")
        yield
    except Exception as e:
        logger.critical(f"服务启动失败: {str(e)}")
        raise
    finally:
        if face_service:
            face_service.close_connections()
        logger.info("人脸识别服务已关闭")


app = FastAPI(
    title="人脸识别API服务",
    description="基于InsightFace和Milvus的面向对象人脸识别服务",
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


def get_face_service() -> FaceRecognitionService:
    if face_service is None:
        raise HTTPException(status_code=500, detail="服务未初始化完成")
    return face_service

@app.get("/")
async def root():
    return {
        "service": "Face Recognition API",
        "version": "1.0.0",
        "endpoints": {
            "health_check": "/health",
            "search_face": "/search_face",
            "docs": "/docs"
        }
    }


@app.get("/health")
async def health_check(service: FaceRecognitionService = Depends(get_face_service)):
    """服务健康检查接口"""
    return {
        "status": "running",
        "milvus_connected": service.cli is not None,
        "rabbitmq_connected": service.rabbitmq_connected,
        "model_loaded": service.face_app is not None
    }


@app.post("/search_face")
async def search_face(
        image: UploadFile = File(...),
        service: FaceRecognitionService = Depends(get_face_service)
):
    try:
        contents = await image.read()
        nparr = np.frombuffer(contents, np.uint8)
        frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        height, width = frame.shape[:2]

        faces = service.face_app.get(frame)
        if not faces:
            return {
                "status": "success",
                "search_results": [],
                "query_face_detected": False,
                "message": "未检测到人脸"
            }

        results = []
        detection_data = {
            'faces': [],
            'frame_shape': [width, height],
            'timestamp': datetime.datetime.now().isoformat()
        }

        for face in faces:
            embedding = face.embedding.tolist()
            search_results = await service.search_similar_faces(embedding, limit=1)
            best_match = search_results[0] if search_results else None

            face_info = {
                'bbox': face.bbox.astype(int).tolist(),
                'key': best_match['key'] if best_match else 'Unknown',
                'confidence': best_match['confidence'] if best_match else 0,
                'distance': best_match['distance'] if best_match else 1.0
            }

            if face_info['key'] != "Unknown":
                detection_data['faces'].append(face_info)

            results.append({
                "bbox": face_info['bbox'],
                "best_match": best_match,
                "distance_threshold": service.settings.similarity_threshold
            })


        if (service.rabbitmq_connected and
            detection_data and
            isinstance(detection_data, dict) and
            detection_data.get('faces') and
            len(detection_data['faces']) > 0):

            service.send_to_rabbitmq({
                "type": "face_detection",
                "timestamp": detection_data['timestamp'],
                "data": {
                    "faces": detection_data['faces'],
                    "frame_shape": detection_data['frame_shape']
                }
            })

        return {
            "status": "success",
            "search_results": results,
            "query_face_detected": True,
            "face_count": len(faces),
            "known_face_count": len([f for f in detection_data['faces'] if f['key'] != 'Unknown']),
            "unknown_face_count": len([f for f in detection_data['faces'] if f['key'] == 'Unknown']),
            "distance_threshold": service.settings.similarity_threshold,
            "timestamp": detection_data['timestamp']
        }

    except Exception as e:
        logger.error(f"人脸搜索失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"人脸搜索失败: {str(e)}")


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8050)
